import os
import pandas as pd
import numpy as np

from utils.draw_pic import plot_single_model_train_and_test, evaluate_model_on_data
from sklearn.metrics import (
    roc_curve, accuracy_score, roc_auc_score,
)


class BaseModelUtil:
    def __init__(self, dp, X, split_random_state, save_path, data_type,model_name,model_random_state):
        self.model_test_result = None
        self.clf = None

        self.dp = dp
        self.X = X
        self.split_random_state = split_random_state
        self.model_random_state = model_random_state
        self.save_path = save_path
        self.data_type = data_type
        self.model_name = model_name

    def run(self):
        if not os.path.exists(self.save_path):
            os.makedirs(self.save_path)

        X_train, X_test, y_train, y_test, _, _ = self.dp.split_train_and_test(self.X,self.data_type,flag=True,model_name="signature_model")
        # 防止KNN因为内存不连续
        X_train = np.ascontiguousarray(X_train.values if hasattr(X_train, 'values') else X_train)
        X_test = np.ascontiguousarray(X_test.values if hasattr(X_test, 'values') else X_test)

        clf = self.build_model()
        clf.fit(X_train, y_train)
        self.clf = clf
        y_train_prob = self.clf.predict_proba(X_train)[:, 1]
        y_test_prob = self.clf.predict_proba(X_test)[:, 1]
        y_train_pred = self.clf.predict(X_train)
        y_test_pred = self.clf.predict(X_test)

        # 评估指标
        auc_train = roc_auc_score(y_train, y_train_prob)
        auc_test = roc_auc_score(y_test, y_test_prob)
        acc_train = accuracy_score(y_train, y_train_pred)
        acc_test = accuracy_score(y_test, y_test_pred)
        # self.model_test_result = evaluate_model_on_data(clf, X_test, y_test)
        # # 可视化
        # plot_single_model_train_and_test(clf,X_train,y_train,X_test, y_test, save_path=self.save_path,draw_ci=True,
        #                                  auc_title=f"{self.data_type} {self.model_name} ROC Curve", cm_title=f"{self.data_type} {self.model_name} Confusion Matrix")
        # # 打印输出
        if 'final' in self.save_path:
        #     if 'all_feature' in self.save_path:
        #         print('使用所有特征的情况下')
        #     else:
        #         print('使用三合一的情况下')
            print(f"模型训练集上AUC = {auc_train:.3f}, Acc = {acc_train:.3f}")
            print(f"模型测试集上AUC = {auc_test:.3f}, Acc = {acc_test:.3f}")


    def get_result(self):
        return self.model_test_result
    
    def get_model(self):
        return self.clf
    
    def getSignature(self):
        X_train, X_test, y_train, y_test, id_train, id_test = self.dp.split_train_and_test(self.X,self.data_type,flag=True)
        # 防止KNN因为数据内存不连续
        X_train = np.ascontiguousarray(X_train.values if hasattr(X_train, 'values') else X_train)
        X_test = np.ascontiguousarray(X_test.values if hasattr(X_test, 'values') else X_test)

        train_sig = self.clf.decision_function(X_train)
        test_sig = self.clf.decision_function(X_test)

        train_sig = pd.DataFrame(train_sig, index=id_train, columns=[self.data_type])
        test_sig = pd.DataFrame(test_sig, index=id_test, columns=[self.data_type])
        return train_sig, test_sig

    def build_model(self, **kwargs):
        raise NotImplementedError("You must implement build_model() in subclass.")

    def get_param_grid(self):
        raise NotImplementedError("You must implement get_param_grid() in subclass.")